import os
import operator
import asyncio
from dotenv import load_dotenv
from typing import Annotated, Any, TypedDict

from langchain_core.messages import AnyMessage, HumanMessage, AIMessage
from langchain_core.tools import tool
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_openai import ChatOpenAI

from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode
from langgraph.checkpoint.memory import MemorySaver

# --- 1. 环境与工具准备 (同教程6) ---
load_dotenv()

# a. 分类工具
try:
    risky_search_tool = TavilySearchResults(max_results=2)
except ImportError:
    from langchain_community.tools import DuckDuckGoSearchRun
    risky_search_tool = DuckDuckGoSearchRun()

@tool
def python_calculator(expression: str) -> str:
    """一个可以执行 Python 数学表达式的计算器。"""
    try:
        result = eval(expression, {"__builtins__": None}, {})
        return f"表达式 '{expression}' 的计算结果是: {result}"
    except Exception as e:
        return f"执行表达式 '{expression}' 时出错: {e}"

safe_tools = [python_calculator]
risky_tools = [risky_search_tool]
all_tools = safe_tools + risky_tools

# --- 2. 状态定义 (同教程6) ---
class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], operator.add]
    iteration_count: int

# --- 3. 异步节点定义 ---
# a. Agent 节点 (大脑) - 必须是异步
llm = ChatOpenAI(model="qwen-plus-latest", base_url=os.getenv("OPENAI_BASE_URL"))
model_with_tools = llm.bind_tools(all_tools)

async def agent_node(state: AgentState):
    print("\n---AGENT: 思考中...---")
    # LangGraph 会自动处理来自内部 ainvoke 调用的流式令牌
    response = await model_with_tools.ainvoke(state["messages"])
    return {"messages": [response]}

# b. 工具节点 - 无需改变
safe_tool_node = ToolNode(safe_tools)
risky_tool_node = ToolNode(risky_tools)

# c. 智能路由节点 - 必须是异步
MAX_ITERATIONS = 5
async def router_node(state: AgentState) -> str:
    print("---ROUTER: 决策中...---")
    
    iteration_count = state.get("iteration_count", 0) + 1
    if iteration_count > MAX_ITERATIONS:
        print(f"已达到最大循环次数 {MAX_ITERATIONS}，流程终止。")
        return END

    last_message = state["messages"][-1]
    if not isinstance(last_message, AIMessage) or not last_message.tool_calls:
        return END

    tool_name = last_message.tool_calls[0]["name"]
    print(f"ROUTER: 计划调用工具 '{tool_name}'")
    
    if tool_name in [t.name for t in safe_tools]:
        print("ROUTER: 路由到 -> 安全工具 (自动执行)")
        return "safe_tools"
    else:
        print("ROUTER: 路由到 -> 风险工具 (需要审批)")
        return "risky_tools"

# --- 4. 构建图 (同教程6) ---
workflow = StateGraph(AgentState)
workflow.add_node("agent", agent_node)
workflow.add_node("safe_tools", safe_tool_node)
workflow.add_node("risky_tools", risky_tool_node)
workflow.set_entry_point("agent")
workflow.add_conditional_edges(
    "agent", router_node, {"safe_tools": "safe_tools", "risky_tools": "risky_tools", END: END}
)
workflow.add_edge("safe_tools", "agent")
workflow.add_edge("risky_tools", "agent")

# --- 5. 编译图，精确设置中断点 (同教程6) ---
checkpointer = MemorySaver()
app = workflow.compile(checkpointer=checkpointer, interrupt_before=["risky_tools"])

# --- 6. 最终结合：流式输出与循环审批 ---
async def main():
    question = "搜索 LangChain 和 LangGraph 的 GitHub 星标数，然后计算它们的总和。"
    config = {"configurable": {"thread_id": "user_session_8"}}
    
    # 跟踪下一次运行的输入。第一次是问题，之后恢复执行时是 None。
    inputs_for_next_run = {
        "messages": [HumanMessage(content=question)],
        "iteration_count": 0
    }

    while True:
        # 使用 astream_events 来运行图，并实时处理流式输出
        print("\n--- Running Graph ---")
        async for event in app.astream_events(inputs_for_next_run, config=config, version="v2"):
            kind = event["event"]
            if kind == "on_chat_model_stream":
                content = event["data"]["chunk"].content
                if content:
                    print(content, end="", flush=True)
        
        # 流结束后，将下一次的输入设为 None 以便恢复
        inputs_for_next_run = None
        
        # 获取中断后的状态快照
        snapshot = await app.aget_state(config)
        
        if snapshot.next:
            print("\n\n---流程暂停，等待您的审批---")
            print(f"下一步计划执行的节点: {snapshot.next[0]}")
            
            if snapshot.next[0] == "risky_tools":
                last_message = snapshot.values["messages"][-1]
                print("风险工具调用详情:", last_message.tool_calls)
            
            user_approval = input("您是否批准执行下一步？ (yes/no/exit): ")
            if user_approval.lower() == "yes":
                print("\n---用户批准，继续执行---")
                continue
            else:
                print("\n---用户拒绝或退出，流程终止---")
                break
        else:
            print("\n\n---流程已完成！---")
            final_message = snapshot.values["messages"][-1]
            print("最终答案:")
            final_message.pretty_print()
            break

if __name__ == "__main__":
    asyncio.run(main()) 